#!/usr/bin/env python

# -----------------------------------------------------
# Written by Zhu Lingfeng on 2020/3/11.
# -----------------------------------------------------

import os
from detectron2.data.datasets.register_coco import register_coco_instances

dataset_name = "coco_format_dataset_demo"
category_names = ["cat1", "cat2"]
data_root = "/home/zhulingfeng/Projects/cocoapi/PythonAPI"
data_dir = "test"

SPLITS = {
    dataset_name + "_train": (data_dir, data_dir + "/test_set.json"),
    dataset_name + "_test": (data_dir, data_dir + "/test_set.json"),
}

for split_name, (image_root, json_file) in SPLITS.items():
    json_file = os.path.join(data_root, json_file)
    image_root = os.path.join(data_root, image_root)
    register_coco_instances(split_name, {}, json_file, image_root, category_names)
